"""
PSPNet (Pyramid Scene Parsing Network) 模型实现
使用金字塔池化模块捕获多尺度上下文信息
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import resnet50, ResNet50_Weights


class PyramidPoolingModule(nn.Module):
    """
    金字塔池化模块
    使用不同尺度的池化捕获全局上下文
    """

    def __init__(self, in_channels, pool_sizes=[1, 2, 3, 6]):
        super(PyramidPoolingModule, self).__init__()

        self.pool_sizes = pool_sizes
        out_channels = in_channels // len(pool_sizes)

        self.stages = nn.ModuleList([
            nn.Sequential(
                nn.AdaptiveAvgPool2d(pool_size),
                nn.Conv2d(in_channels, out_channels, 1, bias=False),
                nn.BatchNorm2d(out_channels),
                nn.ReLU(inplace=True)
            )
            for pool_size in pool_sizes
        ])

    def forward(self, x):
        h, w = x.shape[2:]
        pyramids = [x]

        for stage in self.stages:
            pooled = stage(x)
            # 上采样回原始尺寸
            pooled = F.interpolate(pooled, size=(h, w), mode='bilinear', align_corners=True)
            pyramids.append(pooled)

        # 拼接所有金字塔特征
        output = torch.cat(pyramids, dim=1)

        return output


class PSPNet(nn.Module):
    """
    PSPNet模型
    输入：(B, 3, H, W) - 需要3通道以使用ImageNet预训练
    输出：(B, out_channels, H, W)
    """

    def __init__(self, in_channels=3, out_channels=1, pretrained=True):
        """
        Args:
            in_channels: 输入通道数（必须是3）
            out_channels: 输出通道数
            pretrained: 是否使用预训练权重
        """
        super(PSPNet, self).__init__()

        if in_channels != 3:
            print(f"[WARNING] PSPNet 需要3通道输入以使用预训练权重，当前输入通道数: {in_channels}")

        # 使用ResNet-50作为backbone
        if pretrained:
            resnet = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
        else:
            resnet = resnet50(weights=None)

        # 编码器
        self.conv1 = resnet.conv1
        self.bn1 = resnet.bn1
        self.relu = resnet.relu
        self.maxpool = resnet.maxpool

        self.layer1 = resnet.layer1  # 256 channels
        self.layer2 = resnet.layer2  # 512 channels
        self.layer3 = resnet.layer3  # 1024 channels
        self.layer4 = resnet.layer4  # 2048 channels

        # 金字塔池化模块
        self.ppm = PyramidPoolingModule(in_channels=2048, pool_sizes=[1, 2, 3, 6])

        # PPM输出通道数：2048 + 2048//4 * 4 = 2048 + 512 * 4 = 4096
        ppm_out_channels = 2048 + 2048 // 4 * 4

        # 最终卷积层
        self.final_conv = nn.Sequential(
            nn.Conv2d(ppm_out_channels, 512, 3, padding=1, bias=False),
            nn.BatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Dropout(0.1),
            nn.Conv2d(512, out_channels, 1)
        )

    def forward(self, x):
        size = x.shape[2:]

        # 编码器
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)  # 1/32

        # 金字塔池化
        x = self.ppm(x)

        # 最终卷积
        x = self.final_conv(x)

        # 上采样到原始尺寸
        x = F.interpolate(x, size=size, mode='bilinear', align_corners=True)

        return x


if __name__ == '__main__':
    # 测试模型
    model = PSPNet(in_channels=3, out_channels=1, pretrained=False)
    x = torch.randn(2, 3, 256, 256)
    y = model(x)
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {y.shape}")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.2f}M")
